#!/usr/bin/env python3
"""
Setup script for Weather Prediction AI System.
Handles dependency installation and environment setup.
"""

import os
import sys
import subprocess
import platform

def check_python_version():
    """Check if Python version is compatible"""
    print("Checking Python version...")
    if sys.version_info < (3, 8):
        print(f"❌ Python 3.8+ required. Current version: {sys.version_info.major}.{sys.version_info.minor}")
        return False
    print(f"✓ Python {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} detected")
    return True

def install_dependencies():
    """Install required dependencies"""
    print("\nInstalling dependencies...")
    
    # Check if pip is available
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "--version"])
    except subprocess.CalledProcessError:
        print("❌ pip not available. Please install pip first.")
        return False
    
    # Install requirements
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])
        print("✓ Dependencies installed successfully")
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ Failed to install dependencies: {e}")
        return False

def create_directories():
    """Create necessary directories"""
    print("\nCreating directories...")
    
    directories = [
        './data/himawari/',
        './checkpoints/',
        './logs/', 
        './predictions/'
    ]
    
    for directory in directories:
        try:
            os.makedirs(directory, exist_ok=True)
            print(f"✓ Created directory: {directory}")
        except Exception as e:
            print(f"❌ Failed to create directory {directory}: {e}")
            return False
    
    return True

def verify_installation():
    """Verify that the installation was successful"""
    print("\nVerifying installation...")
    
    # Test basic imports
    try:
        import pandas as pd
        import numpy as np
        import yaml
        print("✓ Basic imports successful")
    except ImportError as e:
        print(f"❌ Basic import failed: {e}")
        return False
    
    # Test if we can read config
    try:
        with open('config.yaml', 'r') as f:
            config = yaml.safe_load(f)
        print("✓ Config file readable")
    except Exception as e:
        print(f"❌ Config file error: {e}")
        return False
    
    # Test if we can read sample data
    try:
        if os.path.exists('station_data.csv'):
            df = pd.read_csv('station_data.csv')
            print(f"✓ Sample data loaded: {len(df)} rows")
        else:
            print("⚠ Sample data file not found (this is normal for fresh install)")
    except Exception as e:
        print(f"❌ Data loading error: {e}")
        return False
    
    return True

def main():
    """Main setup function"""
    print("=" * 60)
    print("WEATHER PREDICTION AI SYSTEM - SETUP")
    print("=" * 60)
    
    # Check Python version
    if not check_python_version():
        return 1
    
    # Create directories
    if not create_directories():
        return 1
    
    # Install dependencies
    if not install_dependencies():
        return 1
    
    # Verify installation
    if not verify_installation():
        return 1
    
    print("\n" + "=" * 60)
    print("🎉 SETUP COMPLETED SUCCESSFULLY!")
    print("=" * 60)
    
    print("\nNext steps:")
    print("1. Add your station data to 'station_data.csv'")
    print("2. Add satellite data to './data/himawari/' directory")
    print("3. Train the model: python train.py")
    print("4. Make predictions: python predict.py")
    print("\nFor testing: python test_system.py")
    
    return 0

if __name__ == "__main__":
    sys.exit(main())
